# coding:utf-8

# from sklearn.feature_extraction.text import CountVectorizer
from sklearn import feature
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
# from sklearn.feature_extraction.text import TfidfVectorizer



def get_source_sample_data():
    from sklearn.datasets import fetch_20newsgroups
    news = fetch_20newsgroups(subset='all')
    return news

def get_train_and_test_data(news):

    from sklearn.cross_validation import train_test_split
    X_train_data , X_test_data, y_train_label , y_test_label = \
        train_test_split(news.data,news.target,test_size=0.25,random_state=22)

    return X_train_data, X_test_data, y_train_label, y_test_label

def count_sample_vec(X_train_data,X_test_data):

    count_vec_no_stopwords = CountVectorizer()
    X_count_train_vec = count_vec_no_stopwords.fit_transform(X_train_data)
    X_count_test_vec = count_vec_no_stopwords.transform(X_test_data)

    return X_count_train_vec, X_count_test_vec

def tfidf_sample_vec(X_train_data,X_test_data,stopwords=True):

    if stopwords:
        tfidf_vec_with_stopwords = TfidfVectorizer(analyzer='word',stopwords='english')
        X_tfidf_train_vec = tfidf_vec_with_stopwords.fit_transform(X_train_data)
        X_tfidf_test_vec = tfidf_vec_with_stopwords.transform(X_test_data)
    else:
        tfidf_vec_no_stopwords = TfidfVectorizer()
        X_tfidf_train_vec = tfidf_vec_no_stopwords.fit_transform(X_train_data)
        X_tfidf_test_vec = tfidf_vec_no_stopwords.transform(X_test_data)

    return X_tfidf_train_vec, X_tfidf_test_vec

def nb_train_model(X_train_vec,y_train_label):

    from sklearn.naive_bayes import MultinomialNB
    mnb = MultinomialNB()
    mnb.fit(X_train_vec,y_train_label)

    return mnb
def count_no_stopwords_model(X_train_vec, y_train_label):

    mnb_count = nb_train_model(X_train_vec, y_train_label)
    return mnb_count

def tfidf_no_stopwords_model(X_train_data, X_test_data, y_train_label, y_test_label,news):
    # tfidf_sample_vec
    # nb_train_model
    # print_evaluate_info
    # evaluate_mole

    X_train_vec, X_tfidf_test_vec = tfidf_sample_vec(X_train_data, X_test_data,stopwords=False)

    mnb_tfidf = nb_train_model(X_train_vec, y_train_label)
    y_predict = mnb_tfidf.predict(X_tfidf_test_vec)

    print_evaluate_info("tfidf", stopwords=False)

    evaluate_mole(mnb_tfidf, X_tfidf_test_vec, y_test_label, y_predict, news)

def count_with_stopwords_model(X_train_vec,y_train_label):
    pass

def tfidf_with_stopwords_model(X_train_data,X_test_data,y_train_label,y_test_label,news):

    X_train_vec, X_tfidf_test_vec = tfidf_sample_vec(X_train_data, X_test_data)

    mnb_tfidf = nb_train_model(X_train_vec, y_train_label)
    y_predict = mnb_tfidf.predict(X_tfidf_test_vec)

    print_evaluate_info("tfidf")

    evaluate_mole(mnb_tfidf, X_tfidf_test_vec, y_test_label, y_predict, news)

def evaluate_mole(mnb,X_test_vec,y_test_label,y_predict,news):

    from sklearn.metrics import classification_report
    accuracy = mnb.score(X_test_vec,y_test_label)
    print(accuracy)
    print(classification_report(y_test_label,y_predict,target_names=news.target_names))

def print_evaluate_info(model_type,stopwords=True):

    if stopwords:

        print("{0} with stopwords evaluate info: ".format(model_type))
    else:
        print("{0} with not stopwords evaluate info: ".format(model_type))

def main():

    # get_source_sample_data
    # get_train_and_test_data
    news = get_source_sample_data()
    X_train_data, X_test_data, y_train_label, y_test_label = get_train_and_test_data(news)
    tfidf_with_stopwords_model(X_train_data, X_test_data, y_train_label, y_test_label, news)
    # tfidf no stopwords



    # count no stopwords
        # count_sample_vec
        # count_no_stopwords_model
        # print_evaluate_info
        # evaluate_mole


    # tfidf with stopwords
        # tfidf_sample_vec
        # tfidf_with_stopwords_model
        # print_evaluate_info
        # evaluate_mole


    # count with stopwords
        # count_sample_vec
        # count_with_stopwords_model
        # print_evaluate_info
        # evaluate_mole


if __name__ == '__main__':
    main()